package org.github.manx98.jlibhac.fssystem;

import lombok.Data;
import org.github.manx98.jlibhac.common.math.IByteUtil;
import org.github.manx98.jlibhac.common.math.Uint32;
import org.github.manx98.jlibhac.common.reflect.StructureMap;
import org.github.manx98.jlibhac.fs.ValueSubStorage;
import org.github.manx98.jlibhac.tools.fssystem.CompressedStorage;
import org.github.manx98.jlibhac.util.BitUtil;

import java.io.IOException;
import java.nio.ByteOrder;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

/**
 * @author Manx98
 * @version 1.0
 * @since 2022/10/3 21:01
 */
@Data
public class BucketTree {
    private static final long ExpectedMagic = 0x52544B42; // BKTR
    private static final int MaxVersion = 1;

    private static final int NodeSizeMin = 1024;
    private static final int NodeSizeMax = 1024 * 512;
    public static final int headerSize = IByteUtil.sizeOf(NodeHeader.class);
    private final OffsetCache offsetCache;
    private NodeBuffer nodeL1;
    private ValueSubStorage nodeStorage;
    private ValueSubStorage entryStorage;
    private int nodeSize;
    private int entrySize;
    private int entryCount;
    private int offsetCount;
    private int entrySetCount;

    public void initialize(ValueSubStorage nodeStorage, ValueSubStorage entryStorage, int nodeSize, int entrySize, int entryCount) throws IOException {
        if (entryCount <= 0) {
            throw new IllegalArgumentException("entryCount must be greater than zero");
        }
        byte[] buffer = new byte[nodeSize];
        nodeStorage.read(0, buffer);
        nodeL1 = new NodeBuffer(buffer);
        nodeL1.header.verify(0, nodeSize, 8);
        // Validate offsets.
        int offsetCount = getOffsetCount(nodeSize);
        int entrySetCount = getEntrySetCount(nodeSize, entrySize, entryCount);
        BucketTreeNode node = nodeL1.getNode();
        long startOffset;
        if (offsetCount < entrySetCount && node.getCount() < offsetCount) {
            startOffset = node.getL2BeginOffset();
        } else {
            startOffset = node.getBeginOffset();
        }

        long endOffset = node.getEndOffset();

        if (startOffset < 0 || startOffset > node.getBeginOffset() || startOffset >= endOffset)
            throw new IllegalArgumentException("InvalidBucketTreeEntryOffset: " + startOffset);
        this.nodeStorage = nodeStorage;
        this.entryStorage = entryStorage;
        this.nodeSize = nodeSize;
        this.entrySize = entrySize;
        this.entryCount = entryCount;
        this.offsetCount = offsetCount;
        this.entrySetCount = entrySetCount;
        this.offsetCache.IsInitialized = true;
        this.offsetCache.Offsets.StartOffset = startOffset;
        this.offsetCache.Offsets.EndOffset = endOffset;
    }

    public Offsets getOffsets() {
        return offsetCache.Offsets;
    }

    public void find(Visitor visitor, long virtualAddress) throws IOException {
        if (virtualAddress < 0)
            throw new IllegalArgumentException(("InvalidOffset: " + virtualAddress));

        if (isEmpty())
            throw new IllegalArgumentException("OutOfRange: " + virtualAddress);

        Offsets offsets = getOffsets();
        visitor.initialize(this, offsets);
        visitor.find(virtualAddress);
    }

    private boolean isEmpty() {
        return entrySize == 0;
    }

    public static class BucketTreeNode {

        private final NodeBuffer nodeBuffer;

        public BucketTreeNode(NodeBuffer nodeBuffer) {
            this.nodeBuffer = nodeBuffer;
        }

        public int getCount() {
            return nodeBuffer.header.EntryCount;
        }

        public long getL2BeginOffset() {
            return IByteUtil.bytesToLong(nodeBuffer.buffer, getCount() * 8, ByteOrder.LITTLE_ENDIAN);
        }

        public long getBeginOffset() {
            return IByteUtil.bytesToLong(nodeBuffer.buffer, 0, ByteOrder.LITTLE_ENDIAN);
        }

        public long getEndOffset() {
            return nodeBuffer.header.OffsetEnd;
        }

        public byte[] getArray() {
            return nodeBuffer.buffer;
        }
    }

    public static class NodeBuffer {
        private final byte[] buffer;
        NodeHeader header = new NodeHeader();

        public NodeBuffer(byte[] buffer) {
            this.buffer = buffer;
            IByteUtil.mapToStruct(header, buffer, 0);
        }

        public BucketTreeNode getNode() {
            return new BucketTreeNode(this);
        }
    }

    public static class OffsetCache {
        public OffsetCache() {
            Mutex = new ReentrantLock();
            IsInitialized = false;
            Offsets.StartOffset = -1;
            Offsets.EndOffset = -1;
        }

        public boolean IsInitialized;
        public Offsets Offsets = new Offsets();
        public Lock Mutex;
    }

    public static class Offsets {
        public long StartOffset;
        public long EndOffset;

        public boolean isInclude(long offset) {
            return StartOffset <= offset && offset < EndOffset;
        }

        public boolean isInclude(long offset, long size) {
            return size > 0 && StartOffset <= offset && size <= EndOffset - offset;
        }
    }

    public BucketTree() {
        this.offsetCache = new OffsetCache();
    }

    public static int queryNodeStorageSize(int nodeSize, int entrySize, int entryCount) {
        if (entryCount <= 0)
            return 0;
        return getEntrySetCount(nodeSize, entrySize, entryCount) * nodeSize;
    }

    private static int getEntrySetCount(int nodeSize, int entrySize, int entryCount) {
        int entryCountPerNode = getEntryCount(nodeSize, entrySize);
        return (int) BitUtil.divideUp(entryCount, entryCountPerNode);
    }

    private static int getEntryCount(int nodeSize, int entrySize) {
        return (nodeSize - headerSize) / entrySize;
    }

    public static int queryEntryStorageSize(int nodeSize, int entrySize, int entryCount) {
        if (entryCount <= 0)
            return 0;

        return (1 + getNodeL2Count(nodeSize, entrySize, entryCount)) * nodeSize;
    }

    private static int getNodeL2Count(int nodeSize, int entrySize, int entryCount) {
        int offsetCountPerNode = getOffsetCount(nodeSize);
        int entrySetCount = getEntrySetCount(nodeSize, entrySize, entryCount);

        if (entrySetCount <= offsetCountPerNode)
            return 0;

        int nodeL2Count = (int) BitUtil.divideUp(entrySetCount, offsetCountPerNode);
        if (nodeL2Count <= offsetCountPerNode) {
            throw new IllegalArgumentException("NodeL2Count " + nodeL2Count + " is out of range, must be between 0 and " + offsetCountPerNode);
        }

        return (int) BitUtil.divideUp(entrySetCount - (offsetCountPerNode - (nodeL2Count - 1)), offsetCountPerNode);
    }

    private static int getOffsetCount(int nodeSize) {
        return (nodeSize - headerSize) / 8;
    }

    public static class Header {
        @StructureMap(offset = 0)
        public Uint32 Magic;
        @StructureMap(offset = 4)
        public Uint32 Version;
        @StructureMap(offset = 8)
        public int EntryCount;

        public void verify() {
            if (Magic.longValue() != ExpectedMagic)
                throw new IllegalArgumentException("Invalid Bucket Tree Signature: " + Magic.longValue());

            if (EntryCount < 0)
                throw new IllegalArgumentException("Invalid Bucket Tree Entry Count: " + EntryCount);

            if (Version.intValue() > MaxVersion)
                throw new IllegalArgumentException("Invalid Bucket Tree Version: " + Version);
        }
    }

    public static class Visitor {
        private BucketTree tree;
        public byte[] entry;
        private Offsets treeOffsets;
        private int entryIndex;
        private int entrySetCount;
        private EntrySetHeader entrySet;
        private final static int entrySetHeaderSize = IByteUtil.sizeOf(EntrySetHeader.class);

        public Visitor() {
            this.tree = null;
            this.treeOffsets = new Offsets();
            this.entryIndex = -1;
            this.entrySetCount = 0;
            this.entrySet = new EntrySetHeader();
        }

        public void initialize(BucketTree tree, Offsets offsets) {
            if (entry == null) {
                entry = new byte[tree.entrySize];
                this.tree = tree;
                treeOffsets = offsets;
            }
        }

        public void find(long virtualAddress) throws IOException {
            // Get the L1 node.
            BucketTreeNode nodeL1 = tree.nodeL1.getNode();

            if (virtualAddress >= nodeL1.getEndOffset())
                throw new IllegalArgumentException("OutOfRange: " + virtualAddress);

            int entrySetIndex;

            if (tree.isExistOffsetL2OnL1() && virtualAddress < nodeL1.getBeginOffset()) {
                // The portion of the L2 offsets containing our target offset is stored in the L1 node
                byte[] offsets = IByteUtil.sliceLong(nodeL1.getArray(), nodeL1.getCount());

                // Find the index of the entry containing the requested offset.
                // If the value is not found, BinarySearch will return the bitwise complement of the
                // index of the first element that is larger than the value.
                // The offsets are the start offsets of each entry, so subtracting 1 from the index of
                // the next-largest value will get us the index of the entry containing the offset.
                int index = IByteUtil.binarySearchLong(offsets, virtualAddress);
                if (index < 0) index = (~index) - 1;

                // If the requested offset comes before the first offset in the list, "index" will be -1.
                if (index < 0)
                    throw new IllegalArgumentException("OutOfRange: " + index);

                entrySetIndex = index;
            } else {
                byte[] offsets = IByteUtil.sliceLong(nodeL1.getArray(), 0, nodeL1.getCount());
                int index = IByteUtil.binarySearchLong(offsets, virtualAddress);
                if (index < 0) index = (~index) - 1;

                if (index < 0)
                    throw new IllegalArgumentException("OutOfRange: " + index);

                if (tree.isExistL2()) {
                    if (index >= tree.offsetCount)
                        throw new IllegalArgumentException("InvalidBucketTreeNodeOffset: " + index);

                    entrySetIndex = findEntrySet(virtualAddress, index);
                } else {
                    entrySetIndex = index;
                }
            }

            // Validate the entry set index.
            if (entrySetIndex < 0 || entrySetIndex >= tree.entrySetCount)
                throw new IllegalArgumentException("InvalidBucketTreeNodeOffset:" + entrySetIndex);

            // Find the entry.
            findEntry(virtualAddress, entrySetIndex);
            // Set count.
            this.entrySetCount = tree.entrySetCount;
        }

        private void findEntry(long virtualAddress, int entrySetIndex) throws IOException {
            findEntryWithoutBuffer(virtualAddress, entrySetIndex);
        }

        private void findEntryWithoutBuffer(long virtualAddress, int entrySetIndex) throws IOException {
            // Calculate entry set extents.
            long entrySize = tree.entrySize;
            long entrySetSize = tree.nodeSize;
            long entrySetOffset = entrySetSize * entrySetIndex;
            ValueSubStorage storage = tree.entryStorage;

            // Read and validate the entry set.
            EntrySetHeader entrySet = new EntrySetHeader();
            byte[] bytes = new byte[BucketTree.headerSize];
            storage.read(entrySetOffset, bytes);
            IByteUtil.mapToStruct(entrySet, bytes, 0);

            entrySet.getHeader().verify(entrySetIndex, entrySetSize, entrySize);

            // Create the node, and find.
            var node = new StorageNode(entrySetOffset, entrySize, entrySet.Info.Count);
            node.find(storage, virtualAddress);

            if (node.getIndex() < 0)
                throw new IllegalArgumentException("InvalidBucketTreeVirtualOffset:" + node.getIndex());

            // Copy the data into entry.
            this.entryIndex = -1;
            int entryIndex = node.getIndex();
            long entryOffset = getBucketTreeEntryOffset(entrySetOffset, entrySize, entryIndex);

            storage.read(entryOffset, entry);

            // Set our entry set/index.
            this.entrySet = entrySet;
            this.entryIndex = entryIndex;
        }

        private int findEntrySet(long virtualAddress, int entrySetIndex) throws IOException {
            return findEntrySetWithoutBuffer(virtualAddress, entrySetIndex);
        }

        private int findEntrySetWithoutBuffer(long virtualAddress, int nodeIndex) throws IOException {
            // Calculate node extents.
            long nodeSize = tree.nodeSize;
            long nodeOffset = nodeSize * (nodeIndex + 1);
            ValueSubStorage storage = tree.nodeStorage;

            // Read and validate the header.
            storage.setPosition(nodeOffset);
            byte[] bytes = new byte[BucketTree.headerSize];
            storage.read(nodeOffset, bytes);
            NodeHeader header = new NodeHeader();
            IByteUtil.mapToStruct(header, bytes, 0);
            header.verify(nodeIndex, nodeSize, 8);
            // Create the node, and find.
            var node = new StorageNode(nodeOffset, 8, header.EntryCount);
            node.find(storage, virtualAddress);

            if (node.getIndex() < 0)
                throw new IllegalArgumentException("InvalidBucketTreeVirtualOffset:" + node.getIndex());

            // Return the index.
            return tree.getEntrySetIndex(header.Index, node.getIndex());
        }

        private long getBucketTreeEntryOffset(long entrySetOffset, long entrySize, int entryIndex) {
            return entrySetOffset + BucketTree.headerSize + entryIndex * entrySize;
        }

        public CompressedStorage.Entry getEntry() {
            CompressedStorage.Entry e = new CompressedStorage.Entry();
            IByteUtil.mapToStruct(e, entry, 0);
            return e;
        }

        public boolean canMoveNext() {
            return isValid() && (entryIndex + 1 < entrySet.Info.Count || entrySet.Info.Index + 1 < entrySetCount);
        }

        private boolean isValid() {
            return entryIndex >= 0;
        }

        public void moveNext() throws IOException {
            if (!isValid())
                throw new IllegalArgumentException("OutOfRange!");

            int entryIndex = this.entryIndex + 1;

            // Invalidate our index, and read the header for the next index.
            if (entryIndex == entrySet.Info.Count) {
                int entrySetIndex = entrySet.Info.Index + 1;
                if (entrySetIndex >= entrySetCount)
                    throw new IllegalArgumentException("OutOfRange: " + entryIndex);

                this.entryIndex = -1;

                long end = entrySet.Info.End;

                long entrySetSize = tree.nodeSize;
                long entrySetOffset = entrySetIndex * entrySetSize;

                byte[] bytes = new byte[entrySetHeaderSize];
                tree.entryStorage.read(entrySetOffset, bytes);
                IByteUtil.mapToStruct(this.entry, bytes, 0);
                entrySet.getHeader().verify(entrySetIndex, entrySetSize, tree.entrySize);
                if (entrySet.Info.Start != end || entrySet.Info.Start >= entrySet.Info.End)
                    throw new IllegalArgumentException("InvalidBucketTreeEntrySetOffset: " + entrySet.Info.Start);

                entryIndex = 0;
            } else {
                this.entryIndex = 1;
            }

            // Read the new entry
            long entrySize = tree.entrySize;
            long entryOffset = getBucketTreeEntryOffset(entrySet.Info.Index, tree.nodeSize, entrySize, entryIndex);
            byte[] bytes = new byte[entrySetHeaderSize];
            tree.entryStorage.read(entryOffset, bytes);
            IByteUtil.mapToStruct(this.entry, bytes, 0);

            // Note that we changed index.
            this.entryIndex = entryIndex;
        }

        private long getBucketTreeEntryOffset(int entrySetIndex, int nodeSize, long entrySize, int entryIndex) {
            return getBucketTreeEntryOffset((long) entrySetIndex * nodeSize, entrySize, entryIndex);
        }
    }

    private int getEntrySetIndex(int nodeIndex, int offsetIndex) {
        return (offsetCount - nodeL1.header.EntryCount) + (offsetCount * nodeIndex) + offsetIndex;
    }

    @Data
    public static class StorageNode {
        private final Offset start;
        private final int count;
        private int index;

        public StorageNode(long offset, long size, int count) {
            this.start = new Offset(offset + BucketTree.headerSize, (int) size);
            this.count = count;
            this.index = -1;
        }

        public void find(ValueSubStorage storage, long virtualAddress) throws IOException {
            int end = count;
            Offset pos = start;

            while (end > 0) {
                int half = end / 2;
                Offset mid = pos.plus(half);

                long offset = 0;
                storage.setPosition(mid.get());
                offset = storage.readLong();
                if (offset <= virtualAddress) {
                    pos = mid.plus(1);
                    end -= half + 1;
                } else {
                    end = half;
                }
            }
            this.index = (int) (pos.reduction(start)) - 1;
        }
    }

    @Data
    public static class Offset {
        private final long offset;
        private final int stride;

        public Offset(long offset, int stride) {
            this.offset = offset;
            this.stride = stride;
        }

        public Offset plusPlus() {
            return plus(1);
        }

        public static Offset reductionReduction(Offset left) {
            return left.reduction(1);
        }

        public Offset plus(long right) {
            return new Offset(offset + right * stride, stride);
        }

        public Offset reduction(long right) {
            return new Offset(offset - right * stride, stride);
        }

        public long reduction(Offset right) {
            return (offset - right.offset) / stride;
        }

        public boolean equal(Offset right) {
            return offset == right.offset;
        }

        public boolean notEqual(Offset right) {
            return offset != right.offset;
        }

        public long get() {
            return offset;
        }
    }

    private boolean isExistOffsetL2OnL1() {
        return isExistL2() && nodeL1.header.EntryCount < offsetCount;
    }

    private boolean isExistL2() {
        return offsetCount < entrySetCount;
    }

    @Data
    public static class EntrySetHeader {
        @StructureMap(offset = 0)
        public NodeHeader Header;
        @StructureMap(offset = 0)
        public EntrySetInfo Info;
    }

    @Data
    private static class EntrySetInfo {
        @StructureMap(offset = 0)
        public int Index;
        @StructureMap(offset = 4)
        public int Count;
        @StructureMap(offset = 8)
        public long End;
        @StructureMap(offset = 16)
        public long Start;
    }
}
